
/*
* This file is part of the openHiTLS project.
*
* openHiTLS is licensed under the Mulan PSL v2.
* You can use this software according to the terms and conditions of the Mulan PSL v2.
* You may obtain a copy of Mulan PSL v2 at:
*
*     http://license.coscl.org.cn/MulanPSL2
*
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
* EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
* MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
* See the Mulan PSL v2 for more details.
*/

#include <string.h>
#include <stdio.h>
#include "api.h"
#include "rand.h"
#include "fips202.h"

//generate keypair
int crypto_kem_keypair(unsigned char *pk, unsigned char *sk)
{
	//call the key generation algorithm of pke
	crypto_encrypt_keypair(pk, sk);
	return 0;
}
int crypto_kem_enc( unsigned char *ct, unsigned char *ss, const unsigned char *pk)
{
	kem_enc_fo(pk,ss,ct);
	return 0;
}
int crypto_kem_dec( unsigned char *ss, const unsigned char *ct, const unsigned char *sk)
{
	const unsigned char *pk;
	pk=sk+SK_LEN;
	kem_dec_fo(pk,sk,ct,ss);
	return 0;
}
// fo encryption for cca security 
int kem_enc_fo(const unsigned char *pk, unsigned char *k, unsigned char *c)
{
	unsigned char buf[MESSAGE_LEN+CIPHER_LEN],seed[SEED_LEN],seed_buf[MESSAGE_LEN+SEED_LEN];
	unsigned long long clen;

	//check parameter
	if(pk==NULL || k==NULL || c==NULL)
	{
		return -1;
	}

	//generate random message m, stored in buf
	random_bytes(buf,MESSAGE_LEN);
	//compute seed=hash(m|pk), add pk for multi key attack protection
	memcpy(seed_buf,buf,MESSAGE_LEN);
	memcpy(seed_buf+MESSAGE_LEN,pk,SEED_LEN);
	sha3_256(seed,seed_buf,MESSAGE_LEN+SEED_LEN);
	//encrypt m with seed
	original_pke_enc_seed(pk,buf,MESSAGE_LEN,c,&clen,seed);
	
	//compute k=hash(m|c)
	memcpy(buf+MESSAGE_LEN,c,CIPHER_LEN);
	sha3_256(k,buf,MESSAGE_LEN+CIPHER_LEN);
	
	return 0;
}

// fo encryption for cca security with seed
int kem_enc_fo_seed(const unsigned char *pk, unsigned char *k, unsigned char *c, unsigned char *seed)
{
	unsigned char buf[MESSAGE_LEN+CIPHER_LEN],local_seed[SEED_LEN],seed_buf[MESSAGE_LEN+SEED_LEN];
	unsigned long long clen;

	
	//check parameter
	if(pk==NULL || k==NULL || c==NULL)
	{
		return -1;
	}

	//generate random message m, stored in buf
	keccak_state state;
    shake256_absorb_once(&state, seed, SEED_LEN);
	shake256_squeeze(buf, MESSAGE_LEN, &state);
	//compute loacal_seed=hash(m|pk), add pk for multi key attack protection
	memcpy(seed_buf,buf,MESSAGE_LEN);
	memcpy(seed_buf+MESSAGE_LEN,pk,SEED_LEN);
	sha3_256(local_seed,seed_buf,MESSAGE_LEN+SEED_LEN);
	//encrypt m with local_seed
	original_pke_enc_seed(pk,buf,MESSAGE_LEN,c,&clen,local_seed);
	
	//compute k=hash(m|c)
	memcpy(buf+MESSAGE_LEN,c,CIPHER_LEN);
	sha3_256(k,buf,MESSAGE_LEN+CIPHER_LEN);
	
	return 0;
}

// decrypt of fo mode
int kem_dec_fo(const unsigned char *pk, const unsigned char *sk, const unsigned char *c, unsigned char *k)
{
	unsigned char buf[MESSAGE_LEN+CIPHER_LEN],seed[SEED_LEN],seed_buf[MESSAGE_LEN+SEED_LEN];
	unsigned long long mlen,clen;
	unsigned char c_v[CIPHER_LEN];
	
	//check parameter
	if(pk==NULL || sk==NULL || k==NULL || c==NULL)
	{
		return -1;
	}
	
	//compute m from c
	original_pke_dec(sk,c,CIPHER_LEN, buf,&mlen);
	//compte k=hash(m|c)
	memcpy(buf+MESSAGE_LEN,c,CIPHER_LEN);
	sha3_256(k,buf,MESSAGE_LEN+CIPHER_LEN);
	//re-encryption with seed=hash(m|pk), add pk for multi key attack protection
	memcpy(seed_buf,buf,MESSAGE_LEN);
	memcpy(seed_buf+MESSAGE_LEN,pk,SEED_LEN);
	sha3_256(seed,seed_buf,MESSAGE_LEN+SEED_LEN);
	original_pke_enc_seed(pk,buf,MESSAGE_LEN,c_v,&clen,seed);
	
	//verify
	if(memcmp(c,c_v,CIPHER_LEN)!=0)
	{
		//k=hash(hash(sk)|c)
		sha3_256(buf,sk,SK_LEN);
		memcpy(buf+MESSAGE_LEN,c,CIPHER_LEN);
		sha3_256(k,buf,MESSAGE_LEN+CIPHER_LEN);
	}
	
	return 0;
}

